import random
from http import HTTPStatus
from dashscope import Generation  # 建议dashscope SDK 的版本 >= 1.14.0
import dashscope
import os


class QwenModel:
    dashscope.api_key = os.environ.get('DASHSCOPE_API_KEY')

    def __init__(self, model='qwen-turbo'):  # qwen-turbo, qwen-max, qwen-plus, qwen1.5-110b-chat, qwen-long, qwen1.5-72b-chat, qwen1.5-32b-chat
        self.model = model

    def chat(self, prompt):
        messages = [{'role': 'system', 'content': 'You are a helpful assistant.'},
                    {'role': 'user', 'content': prompt}]
        response = Generation.call(model=self.model,
                                messages=messages,
                                # 设置随机数种子seed，如果没有设置，则随机数种子默认为1234
                                seed=random.randint(1, 10000),
                                # 将输出设置为"message"格式
                                result_format='message')
        if response.status_code == HTTPStatus.OK:
            # print(response)
            return response.output.choices[0].message.content
        else:
            print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
                response.request_id, response.status_code,
                response.code, response.message
            ))
            return None


if __name__ == '__main__':
    # call_with_messages()
    qwen = QwenModel(model='llama3-70b-instruct')
    prompt_prefix = '按行判断下面每行文字是什么语言。然后将它翻译成英语和简体中文。最后按csv格式输出结果：第一列为原文，第二列为2字母的语言代码，第三列为英语结果，第四列为中文结果。如果原文或译文包含逗号或者引号，按照csv的标准方式进行转义处理。结果只需要输出一遍。不要输出csv表头。如果文字无法翻译，就将这一行舍弃。'
    prompt_content = 'una mente libera\nuncinetto che conta\nkit gilet jaune et triangle de signalisation voiture\ncavo antifurto casco\ntete de débroussailleuse husqvarna\nportachiavi juventus ufficiale 2024\nmanette ps5 grise\n2 bedroom tiny house\npersonalised stanley cup'
    result = qwen.chat(prompt_prefix + '\n' + prompt_content)
    print(result)
